# coding=utf-8

import hiai
from hiai.nn_tensor_lib import DataType
from utils import jpegHandler
import os
import numpy as np
import time
import cv2 as cv
label_2_num_txt = './coco_label_2_num.txt'
with open(label_2_num_txt, 'r') as f:
    lines = f.readlines()
class_names = []
for line in lines:
    class_names.append(line.strip())

num_class = len(class_names)
index_to_class = dict(zip(range(num_class), class_names))


omFileName = './models/yolov1_tiny_deploy.om'
srcFileDir = './ImageNetRaw/'
dstFileDir = './Result/yolov1_tiny_deploy'
kCategoryIndex = 2

if not os.path.exists(dstFileDir):
    os.mkdir(dstFileDir)

def CreateGraph(model, modelInWidth, modelInHeight, dvppInWidth, dvppInHeight):
    myGraph = hiai.hiai._global_default_graph_stack.get_default_graph()
    if myGraph is None:
        print 'get defaule graph failed'
        return None
    print 'dvppwidth %d, dvppheight %d' % (dvppInWidth, dvppInHeight)

    cropConfig = hiai.CropConfig(0, 0, dvppInWidth, dvppInHeight)
    print 'cropConfig ', cropConfig
    resizeConfig = hiai.ResizeConfig(modelInWidth, modelInHeight)
    print 'resizeConfig ', resizeConfig

    nntensorList = hiai.NNTensorList()
    print 'nntensorList', nntensorList

    resultCrop = hiai.crop(nntensorList, cropConfig)
    print 'resultCrop', resultCrop

    resultResize = hiai.resize(resultCrop, resizeConfig)
    print 'resultResize', resultResize

    resultInference = hiai.inference(resultResize, model, None)
    print 'resultInference', resultInference

    if (hiai.HiaiPythonStatust.HIAI_PYTHON_OK == myGraph.create_graph()):
        print 'create graph ok !!!!'
        return myGraph
    else:
        print 'create graph failed, please check Davinc log.'
        return None


def CreateGraphWithoutDVPP(model):
    print model
    myGraph = hiai.hiai._global_default_graph_stack.get_default_graph()
    print myGraph
    if myGraph is None:
        print 'get defaule graph failed'
        return None

    nntensorList = hiai.NNTensorList()
    print nntensorList

    resultInference = hiai.inference(nntensorList, model, None)
    print nntensorList
    print hiai.HiaiPythonStatust.HIAI_PYTHON_OK
    # print myGraph.create_graph()

    if (hiai.HiaiPythonStatust.HIAI_PYTHON_OK == myGraph.create_graph()):
        print 'create graph ok !!!!'
        return myGraph
    else:
        print 'create graph failed, please check Davinc log.'
        return None


def GraphInference(graphHandle, inputTensorList):
    if not isinstance(graphHandle, hiai.Graph):
        print "graphHandle is not Graph object"
        return None

    resultList = graphHandle.proc(inputTensorList)
    return resultList

def PostProcess(resultList, dstFilePath, fileName):
    confidence_thresh = 0.5
    if resultList is not None:
        tensor_bbox = np.reshape(resultList[0], (30,49))
        tensor_num = np.reshape(resultList[1], (1))
        img = cv.imread(fileName)

        img_rows, img_cols, img_channel = img.shape


        idx = np.where(tensor_bbox[:, 2] >= confidence_thresh)
        bbox_keep = tensor_bbox[idx, :]
        print('bboxes to keep', bbox_keep)
        for i in range(len(bbox_keep[0])):
            lt_x = int(img_cols * bbox_keep[0][i][3])
            lt_y = int(img_rows * bbox_keep[0][i][4])
            rb_x = int(img_cols * bbox_keep[0][i][5])
            rb_y = int(img_rows * bbox_keep[0][i][6])
            score = bbox_keep[0][i][2]
            cls = int(bbox_keep[0][i][1])
            cv.rectangle(img, (lt_x, lt_y), (rb_x, rb_y), (0, 0, 255), 4)
            cv.putText(img, "class:%s,score:%.3f" % (index_to_class[cls], score), (lt_x, lt_y),
                       cv.FONT_HERSHEY_COMPLEX, 0.5, (0, 255, 255), 1)

        _, filename = os.path.split(fileName)
        cv.imwrite("%s%s" % (dstFilePath, filename), img)
        return None
    else:
        print 'graph inference failed '
        return None



def main():
    # capture_flag = 1  # use camera or not
    inferenceModel = hiai.AIModelDescription('vgg-ssd', omFileName)
    print omFileName
    print inferenceModel
    # we will resize the jpeg to 896*608 to meet faster-rcnn requirement via opencv,
    # so DVPP resizing is not needed
    myGraph = CreateGraphWithoutDVPP(inferenceModel)
    if myGraph is None:
        print "CreateGraph failed"
        return None

    # in this sample demo, the faster-rcnn  model requires 896*608 images
    dvppInWidth = 512
    dvppInHeight = 448

    start = time.time()
    index = 0

    pathDir = sorted(os.listdir(srcFileDir))
    for allDir in pathDir:
        child = os.path.join('%s%s' % (srcFileDir, allDir))
        if (not jpegHandler.is_img(child)):
            print '[info] file : ' + child + ' is not image !'
            continue

        # read the jpeg file and resize it to required w&h, than change it to YUV format.
        input_image = jpegHandler.jpeg2yuv(child, dvppInWidth, dvppInHeight)

        inputImageTensor = hiai.NNTensor(input_image, dvppInWidth, dvppInHeight, 3, 'testImage', DataType.UINT8_T,
                                         dvppInWidth * dvppInHeight * 3 / 2)
        # imageinfo = np.array([dvppInWidth, dvppInHeight, 3]).astype(np.float32)
        # imageinfo = np.reshape(imageinfo, (1, 3))
        # infoTensor = hiai.NNTensor(imageinfo, 1, 3, 1, 'testinfo', DataType.FLOAT32_T, imageinfo.size)

        # datalist = [inputImageTensor, infoTensor]
        datalist = [inputImageTensor]

        nntensorList = hiai.NNTensorList(datalist)
        if not nntensorList:
            print "nntensorList is null"
            break
        resultList = GraphInference(myGraph, nntensorList)

        if resultList is None:
            print "graph inference failed"
            continue
        print resultList[0].shape
        PostProcess(resultList, dstFileDir, child)

    end = time.time()
    print 'cost time ' + str((end - start) * 1000) + 'ms'

    hiai.hiai._global_default_graph_stack.get_default_graph().destroy()

    print '-------------------end'


if __name__ == "__main__":
    main()

